import os
from pathlib import Path

import pandas as pd
import glob


def generate_csv(release_file, file_output):
    fold_data = []
    for label_type in os.listdir(release_file):
        label_type_path = os.path.join(release_file, label_type)
        for image in glob.glob(os.path.join(label_type_path, '*/*.png')):
            image_path = os.path.join(*Path(image).parts[-4:])
            # label = 1 if label_type == 'attack' else 0
            label = 0 if label_type == 'attack' else 1
            fold_data.append([image_path, label])
    columns = ["now_path", "target"]
    pd.DataFrame(fold_data, columns=columns).to_csv(file_output, index=False)


def main():
    data_dir = '/home/shaohua/data/heyan/Datasets/Face/casia_256'
    train_release = os.path.join(data_dir, 'train_release')
    test_release = os.path.join(data_dir, 'test_release')
    generate_csv(train_release, '../anno_file/train_normal_sample.csv')
    generate_csv(test_release, '../anno_file/val_normal_sample.csv')


if __name__ == '__main__':
    main()


'''
# old version
def generate_csv(data_dir, file_input, file_output):
    data_file = os.path.join(data_dir, file_input)
    data = open(data_file).readlines()
    fold_data = []
    for line in data:
        line_split = line.split()
        assert len(line_split) == 4
        target = 1 if int(line_split[3]) == 0 else 0
        fold_data.append([line_split[0], target])
    columns = ["path", "target"]
    pd.DataFrame(fold_data, columns=columns).to_csv(file_output, index=False)


def main():
    data_dir = 'D:\\Datasets\\Face_Anti_Spoofing\\CASIA-SURF\\challenge\\phase1'
    generate_csv(data_dir, 'train_list.txt', 'anno_file/train_normal_sample.csv')
    generate_csv(data_dir, 'val_private_list.txt', 'anno_file/val_normal_sample.csv')


if __name__ == '__main__':
    main()


'''